package edu.zju.gis.dbfg.model.tile;

import edu.zju.gis.dbfg.common.constant.Seprator;
import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.Pipeline;
import edu.zju.gis.dbfg.model.base.PipelineBuilder;
import edu.zju.gis.dbfg.model.base.VectorTileBuilder;
import edu.zju.gis.dbfg.model.io.geojson.GeoJsonTileBuilder;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import edu.zju.gis.dbfg.model.util.TileUtil;
import edu.zju.gis.dbfg.model.util.ZLevelInfo;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import org.apache.spark.broadcast.Broadcast;
import org.elasticsearch.spark.rdd.EsSpark;
import org.geotools.geometry.jts.JTS;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.io.WKTReader;
import org.opengis.referencing.FactoryException;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.awt.*;
import java.util.*;
import java.util.List;

/**
 * @author Hu
 * @date 2019/9/4
 * 基于Elasticsearch数据源，实现矢量瓦片构建
 * 目前只支持索引全量数据构建瓦片，不支持数据过滤及字段选择
 * TODO 默认输入输出均为4490，且assume 4490 与 4326 一直，没做坐标系转换
 **/
public class VectorTileClipper {

  private final static Logger logger = LoggerFactory.getLogger(VectorTileClipper.class);

  public static void main(String[] args) throws Exception {

    // Setup environment
    logger.info("Setup Spark Context");
    SparkConf conf = new SparkConf();
    conf.setAppName("vector tile clipper");
    // conf.setMaster("yarn-cluster");
    conf.setMaster("local[4]");
    conf.set("spark.es.resource", "index/type");
    JavaSparkContext jsc = new JavaSparkContext(conf);

    // Setup model arguments
    logger.info("Setup arguments");
    VectorTileClipperArgs mArgs = VectorTileClipperArgs.initArgs(args);
    if (mArgs == null){
      String msg = "init args failed, exit";
      logger.error(msg);
      throw new ModelFailedException(VectorTileClipper.class, "main()", msg, args);
    }

    // TODO 默认输入输出坐标系均为4490，单位为度，暂不支持米制坐标系
    CoordinateReferenceSystem crs = CRS.decode(mArgs.getTargetCRS());
    // 定义金字塔网格参数
    PyramidConfig pyramidConfig = new PyramidConfig.PyramidConfigBuilder()
      .setBaseMapEnv(-180, 180, -90, 90)
      .setCrs(crs)
      .setzLevelRange(mArgs.getZMin(), mArgs.getZMax())
      .build();

//    double[] extent = Arrays.stream(mArgs.getExtent().split(Seprator.COMMA.getValue())).mapToDouble(Double::valueOf).toArray();
//    Envelope envelope = new Envelope(extent[0], extent[1], extent[2], extent[3]);
//    ZLevelInfo[] zLevelInfos = TileUtil.initZLevelInfoPZ(pyramidConfig, envelope);

    // Setup data from datasource
    logger.info("Setup datasource");
    String datasource = mArgs.getInput();
//    if (!datasource.startsWith("es://")) {
//      String msg = String.format("Unvalid datasource %s, only support elasticsearch datastore yet (url schema as 'es://{indexname}')", datasource);
//      logger.error(msg);
//      throw new ModelFailedException(VectorTileClipper.class, "main()", msg, args);
//    }

    // 默认 buffer 为 1
    final Broadcast<Integer> bufferBroad = jsc.broadcast(1);
    final Broadcast<VectorTileClipperArgs> mArgsBroad = jsc.broadcast(mArgs);
    final Broadcast<PyramidConfig> pyramidConfigBroad = jsc.broadcast(pyramidConfig);

    // TODO 如果硬件内存无法容纳所有的数据一起计算，那就在这里进行计算的时候，分批次读取数据并计算
    // TODO 或者，第一层级的数据可以通过更上一层级得数据得到，感觉这个更靠谱一些
    logger.info("Model run");
    for(int z=mArgs.getZMax(); z<=mArgs.getZMin(); z--){
      logger.info(String.format("================== LEVEL %d START ==================", z));
      // Transform data to RDD
      logger.info(String.format("Transform data to RDD from datasource %s", datasource));
      // TODO 将数据源改为从es读取: https://www.elastic.co/guide/en/elasticsearch/hadoop/7.3/spark.html#spark-native
      JavaRDD<GeoItem> geometry = jsc.textFile(datasource).repartition(4).map(new Function<String, GeoItem>() {
        @Override
        public GeoItem call(String s) throws Exception {
          String[] fields = s.split(Seprator.TAB.getValue());
          WKTReader wktReader = new WKTReader();
          Geometry geometry = wktReader.read(fields[fields.length-1]);
          // TODO 测试，暂时先用两个字段
          // TODO 默认第一列为 FieldId
          String fid = fields[0];
          Map<String, Object> prop = new HashMap<>();
          prop.put("CC", fields[2]);
          GeoItem geoItem = new GeoItem();
          geoItem.setFid(fid);
          geoItem.setProp(prop);
          geoItem.setGeometry(geometry);
          return geoItem;
        }
      });

      // JavaPairRDD<String, Map<String, Object>> esRDD = JavaEsSpark.esRDD(jsc, "radio/artists");

      // 每一个图斑单独处理，将其映射到瓦片中
      // key 为 tile ID，value 为对应的Geometry
      JavaPairRDD<TileID, GeoItem> tileGeometry= geometry.flatMapToPair(new PairFlatMapFunction<GeoItem, TileID, GeoItem>() {
        @Override
        public Iterator<Tuple2<TileID, GeoItem>> call(GeoItem geoItem) throws Exception {
          List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();
          Geometry geometry = geoItem.getGeometry();
          // 获取geometry对应的envelope
          ReferencedEnvelope envelope = JTS.toEnvelope(geometry);
          // 获取geometry envelope对应的Tile范围
          ZLevelInfo[] tZLevelInfos = TileUtil.initZLevelInfoPZ(pyramidConfigBroad.getValue(), envelope);
          for(int z=mArgsBroad.getValue().getZMin(); z<=mArgsBroad.getValue().getZMax(); z++){
            ZLevelInfo tZLevelInfo = tZLevelInfos[z-mArgsBroad.getValue().getZMin()];
            // 针对每一层级的不同瓦片，构造图斑处理的pipeline
            int tx_min = tZLevelInfo.getTileRanges()[0];
            int tx_max = tZLevelInfo.getTileRanges()[1];
            int ty_min = tZLevelInfo.getTileRanges()[2];
            int ty_max = tZLevelInfo.getTileRanges()[3];
            for (int tile_x = tx_min; tile_x <= tx_max; tile_x ++) {
              for (int tile_y = ty_min; tile_y <= ty_max; tile_y ++) {
                TileID tileID = new TileID();
                tileID.setX(tile_x);
                tileID.setY(tile_y);
                tileID.setzLevel(z);
                Envelope tileEnvelope = TileUtil.createTileBox(tileID, pyramidConfigBroad.getValue());
                Pipeline pipeline = getPipeline(CRS.decode(mArgsBroad.getValue().getTargetCRS()), bufferBroad.getValue(), tileEnvelope, false);
                // 得到在每一个瓦片中对应的geometry
                Geometry finalGeom = pipeline.execute(geometry);
                if(!finalGeom.isEmpty()){
                  GeoItem finalItem = new GeoItem();
                  finalItem.setGeometry(finalGeom);
                  finalItem.setProp(geoItem.getProp());
                  finalItem.setFid(geoItem.getFid());
                  result.add(new Tuple2<>(tileID, finalItem));
                }
              }
            }
          }
          return result.iterator();
        }
      });

      // 根据瓦片ID group 对应的 geometry，并写出
      JavaRDD<String> result = tileGeometry.groupByKey().map(new Function<Tuple2<TileID, Iterable<GeoItem>>, String>() {
        @Override
        public String call(Tuple2<TileID, Iterable<GeoItem>> tileIDIterableTuple2) throws Exception {
          TileID tileID = tileIDIterableTuple2._1;
          Iterator<GeoItem> geoItemIterator = tileIDIterableTuple2._2.iterator();
          Envelope tileEnvelope = TileUtil.createTileBox(tileID, pyramidConfigBroad.getValue());
          Pipeline pipeline = getPipeline(CRS.decode(mArgsBroad.getValue().getTargetCRS()), bufferBroad.getValue(), tileEnvelope, true);
          VectorTileBuilder vectorTileBuilder = new GeoJsonTileBuilder(mArgsBroad.getValue().getOutDir(), tileID, pyramidConfigBroad.getValue());
          int count = 0;
          while(geoItemIterator.hasNext()){
            GeoItem geoItem = geoItemIterator.next();
            Geometry finalGeom = pipeline.execute(geoItem.getGeometry());
            if (!finalGeom.isEmpty()){
              // 默认 geometryName = Fid
              vectorTileBuilder.addFeature(mArgsBroad.getValue().getLayerName(), geoItem.getFid(), geoItem.getFid(),
                finalGeom, geoItem.getProp());
              count ++;
            }
          }
          vectorTileBuilder.build();
          String msg = String.format("tile[%d, %d, %d] success with total count [%d]", tileID.getzLevel(), tileID.getX(), tileID.getY(), count);
          System.out.println(msg);
          return msg;
        }
      });

      List<String> resultStr = result.collect();
      resultStr.forEach(logger::info);
      logger.info(String.format("Total tile count = %d", resultStr.size()));

      logger.info(String.format("================== LEVEL %d FINISH ==================", z));
    }
    // Finish calculation
    logger.info("Model Finish");
  }

  /**
   * 构建瓦片中每个像素处理的Pipeline
   * 默认瓦片大小：256 * 256
   * OverSampleFactor = 2
   * 默认目标crs与源crs一致
   */
  private static Pipeline getPipeline(CoordinateReferenceSystem sourceCrs, int buffer, Envelope referenceEnvelop, boolean preprocess) throws Exception {
    CoordinateReferenceSystem targetCrs = sourceCrs;
    int mapWidth = 256;
    int mapHeight = 256;
    int overSampleFactor = 2;
    // rectangle of the image: width and height
    Rectangle paintArea = new Rectangle(mapWidth, mapHeight);

    // bounding box - in final map (target) CRS
    ReferencedEnvelope renderingArea = new ReferencedEnvelope(referenceEnvelop.getMinX(), referenceEnvelop.getMaxX(), referenceEnvelop.getMinY(), referenceEnvelop.getMaxY(), targetCrs);
    Pipeline pipeline = null;
    try {
      final PipelineBuilder builder =
        PipelineBuilder.newBuilder(
          renderingArea, paintArea, sourceCrs, overSampleFactor, buffer);

      if (preprocess){
        pipeline =
          builder.preprocess().build();
      } else {
        pipeline =
          builder.transform(false)
            .clip(true, false)
            .simplify(false)
            .collapseCollections()
            .build();
      }
    } catch (FactoryException e) {
      e.printStackTrace();
    }

    return pipeline;
  }

}
